import os
import json
import socket
import hmac
import hashlib
from typing import Dict, Any, Optional
from dotenv import load_dotenv
from loguru import logger

# 加载环境变量
load_dotenv()

class MCPClient:
    """MCP协议客户端，用于与服务器通信"""
    
    def __init__(self, host: Optional[str] = None, port: Optional[int] = None, secret: Optional[str] = None):
        """
        初始化MCP客户端
        
        Args:
            host: 服务器主机名，默认从环境变量获取
            port: 服务器端口，默认从环境变量获取
            secret: 通信密钥，默认从环境变量获取
        """
        self.host = host or os.getenv("MCP_HOST", "localhost")
        self.port = port or int(os.getenv("MCP_PORT", "8080"))
        self.secret = secret or os.getenv("MCP_SECRET", "")
        self.socket = None
        self.connected = False
        
        if not self.secret:
            logger.warning("MCP客户端未配置密钥，将使用不安全的通信")
    
    def connect(self) -> bool:
        """连接到MCP服务器"""
        if self.connected:
            return True
        
        # 重新加载环境变量配置
        load_dotenv(override=True)
        
        # 更新连接参数
        self.host = os.getenv("MCP_HOST", self.host)
        self.port = int(os.getenv("MCP_PORT", str(self.port)))
        self.secret = os.getenv("MCP_SECRET", self.secret)
        
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.host, self.port))
            self.connected = True
            logger.info(f"已连接到MCP服务器 {self.host}:{self.port}")
            return True
        except Exception as e:
            logger.error(f"连接MCP服务器失败: {str(e)}")
            return False
    
    def disconnect(self):
        """断开与服务器的连接"""
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
            self.socket = None
        self.connected = False
        logger.info("已断开与MCP服务器的连接")
    
    def send_request(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        发送请求到服务器并接收响应
        
        Args:
            data: 请求数据
            
        Returns:
            服务器响应数据
        """
        if not self.connected:
            if not self.connect():
                return {"status": "error", "message": "无法连接到服务器"}
        
        # 构建消息
        message = {"data": data}
        
        # 添加签名（如果启用）
        if self.secret:
            message_str = json.dumps(message, sort_keys=True)
            signature = hmac.new(
                self.secret.encode('utf-8'),
                message_str.encode('utf-8'),
                hashlib.sha256
            ).hexdigest()
            message["signature"] = signature
        
        # 编码消息
        message_bytes = json.dumps(message).encode('utf-8')
        
        try:
            # 发送消息长度和数据
            length_bytes = len(message_bytes).to_bytes(4, byteorder='big')
            self.socket.sendall(length_bytes + message_bytes)
            
            # 接收响应长度
            length_data = self.socket.recv(4)
            if not length_data:
                raise ConnectionError("服务器关闭了连接")
            
            response_length = int.from_bytes(length_data, byteorder='big')
            
            # 接收完整响应
            chunks = []
            bytes_received = 0
            while bytes_received < response_length:
                chunk = self.socket.recv(min(response_length - bytes_received, 4096))
                if not chunk:
                    raise ConnectionError("接收响应时连接中断")
                chunks.append(chunk)
                bytes_received += len(chunk)
            
            # 组合完整响应
            response_data = b''.join(chunks)
            response = json.loads(response_data.decode('utf-8'))
            
            # 验证响应签名（如果启用）
            if self.secret and "signature" in response:
                if not self._verify_response(response):
                    logger.warning("响应签名验证失败")
                    return {"status": "error", "message": "响应签名验证失败"}
            
            return response
        
        except ConnectionError as e:
            logger.error(f"与服务器的连接中断: {str(e)}")
            self.disconnect()
            return {"status": "error", "message": f"连接中断: {str(e)}"}
        except Exception as e:
            logger.error(f"发送请求时出错: {str(e)}")
            return {"status": "error", "message": f"请求失败: {str(e)}"}
    
    def _verify_response(self, response: Dict[str, Any]) -> bool:
        """验证响应的签名"""
        if not self.secret:
            return True
        
        signature = response.get("signature")
        if not signature:
            return False
        
        # 移除签名字段后计算响应的HMAC
        response_copy = response.copy()
        response_copy.pop("signature", None)
        response_str = json.dumps(response_copy, sort_keys=True)
        
        # 计算HMAC
        expected_signature = hmac.new(
            self.secret.encode('utf-8'),
            response_str.encode('utf-8'),
            hashlib.sha256
        ).hexdigest()
        
        return hmac.compare_digest(signature, expected_signature)
    
    def execute_query(self, natural_language_query: str) -> Dict[str, Any]:
        """
        执行自然语言查询
        
        Args:
            natural_language_query: 自然语言查询
            
        Returns:
            查询结果
        """
        request = {
            "action": "query",
            "query": natural_language_query
        }
        
        return self.send_request(request)
    
    def get_database_info(self) -> Dict[str, Any]:
        """
        获取数据库信息
        
        Returns:
            数据库信息，包括表结构等
        """
        request = {
            "action": "info"
        }
        
        return self.send_request(request)
    
    def execute_sql(self, sql_query: str) -> Dict[str, Any]:
        """
        执行原始SQL查询
        
        Args:
            sql_query: SQL查询语句
            
        Returns:
            查询结果
        """
        request = {
            "action": "sql",
            "sql": sql_query
        }
        
        return self.send_request(request) 